// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/parquet/reader/nested_structure_decoder.h>

#include <pollux/dwio/common/buffer_util.h>

namespace kumo::pollux::parquet {
    int64_t NestedStructureDecoder::readOffsetsAndNulls(
        const uint8_t *definitionLevels,
        const uint8_t *repetitionLevels,
        int64_t numValues,
        uint8_t maxDefinition,
        uint8_t maxRepeat,
        BufferPtr &offsetsBuffer,
        BufferPtr &lengthsBuffer,
        BufferPtr &nullsBuffer,
        memory::MemoryPool &pool) {
        dwio::common::ensureCapacity<uint8_t>(
            nullsBuffer, bits::nbytes(numValues), &pool);
        dwio::common::ensureCapacity<vector_size_t>(
            offsetsBuffer, numValues + 1, &pool);
        dwio::common::ensureCapacity<vector_size_t>(lengthsBuffer, numValues, &pool);

        auto offsets = offsetsBuffer->asMutable<vector_size_t>();
        auto lengths = lengthsBuffer->asMutable<vector_size_t>();
        auto nulls = nullsBuffer->asMutable<uint64_t>();

        int64_t offset = 0;
        int64_t lastOffset = 0;
        bool wasLastCollectionNull = definitionLevels[0] == (maxDefinition - 1);
        bits::set_null(nulls, 0, wasLastCollectionNull);
        offsets[0] = 0;

        int64_t outputIndex = 1;
        for (int64_t i = 1; i < numValues; ++i) {
            uint8_t definitionLevel = definitionLevels[i];
            uint8_t repetitionLevel = repetitionLevels[i];

            // empty means it belongs to a row that is null in one of its ancestor
            // levels.
            bool isEmpty = definitionLevel < (maxDefinition - 1);
            bool isNull = definitionLevel == (maxDefinition - 1);
            bool isCollectionBegin = (repetitionLevel < maxRepeat) & !isEmpty;
            bool isEntryBegin = (repetitionLevel <= maxRepeat) & !isEmpty;

            offset += isEntryBegin & !wasLastCollectionNull;
            offsets[outputIndex] = offset;
            lengths[outputIndex - 1] = offset - offsets[outputIndex - 1];
            bits::set_null(nulls, outputIndex, isNull);

            // Always update the outputs, but only increase the outputIndex when the
            // current entry is the begin of a new collection, and it's not empty.
            // Benchmark shows skipping non-collection-begin rows is worse than this
            // solution by nearly 2x because of extra branchings added for skipping.
            outputIndex += isCollectionBegin;
            wasLastCollectionNull = isEmpty ? wasLastCollectionNull : isNull;
            lastOffset = isCollectionBegin ? offset : lastOffset;
        }

        offset += !wasLastCollectionNull;
        offsets[outputIndex] = offset;
        lengths[outputIndex - 1] = offset - lastOffset;

        return outputIndex;
    }
} // namespace kumo::pollux::parquet
